import pandas as pd

A = pd.DataFrame([[1, 1, 1], [1, 2, 1], [1, 2, 2],
                  [1, 3, 1], [2, 1, 1]], columns=['id', 'country', 'source'])

A.groupby('id').apply(print)
A.groupby('id')['country'].agg(list)
A.groupby('id').agg(list)['country']
A.groupby('id')['country'].nunique()
A.groupby('id').nunique()

# 闭包函数
def outer(i):
    def inner(x):
        return x+i
    return inner


test = outer(5)
A.groupby('id').apply(test)


# agg函数
def outer_cnt(i):
    def inner(x):
        return x.shape[0]+i
    return inner


A.groupby('id').agg(outer_cnt(1))
A.groupby('id')['country'].agg({'a1': outer_cnt(1)})
A.groupby('id')['country'].agg({'a1': outer_cnt(1), 'a2': outer_cnt(2), 'a3_cnt': 'count'})


def diff(a):
    arr = a.tolist()
    b = pd.Series(arr[1:] + [arr[-1]]) - a
    return b
A.apply(diff, axis=0)
A.groupby('id')['country'].apply(diff, **{'col': 'country'})
A['ord'] = range(A.shape[0])


def desc_shape(a):
    return a.shape


A.groupby('id').apply(desc_shape)


def data_rec(a):
    print(a)
    return a
A.groupby('id').apply(data_rec)


def diff_op(df):
    def diff1(a):
        arr = a.tolist()
        b = pd.Series(arr[1:] + [arr[-1]]) - a
        return b
    df['ord'] = range(df.shape[0])
    df1 = df.sort_values(by='country').reset_index().drop('index', axis=1)
    feas = [a for a in df1.columns if a not in ['ord']]
    df1[feas] = df1[feas].apply(diff1, axis=0)
    return df1


A.groupby('id').apply(diff_op)


def diff_op(df, cols_sp):
    def diff1(a, cols_sp):
        print(cols_sp)
        if a.name not in cols_sp:
            arr = a.tolist()
            b = pd.Series(arr[1:] + [arr[-1]]) - a
        else:
            b = a
        return b
    df1 = df.sort_values(by='country').reset_index().drop('index', axis=1)
    df2 = df1.apply(diff1, axis=0, **{'cols_sp': cols_sp})
    return df2


cols_sp = ['ord']
A.groupby('id').apply(diff_op, cols_sp)